#!/usr/bin/env python

import sys
import rospy
from sensor_msgs.msg import Image
from vision_messages.srv import InstanceMask
import numpy as np
import cv2

cate_map = {
    5: 'mug',
    1: 'bowl',
    0: 'bottle'
}

class Mask_Client():
    def __init__(self):
        rospy.init_node('mask_client')
        rospy.Subscriber("/camera/color/image_raw", Image, self.image_callback)
        self.image = rospy.wait_for_message("/camera/color/image_raw", Image)
        self.depth = self.image

    def mask_client(self):
        rospy.wait_for_service('/perception/mmdseg/solo')
        try:
            get_mask = rospy.ServiceProxy('/perception/mmdseg/solo', InstanceMask)
            result = get_mask(self.image, self.depth)
            cate = result.cate
            return cate
        except rospy.ServiceException as e:            
            rospy.loginfo("Service call failed: %s"%e)

    def image_callback(self, data):
        self.image = data
        self.depth = data

if __name__ == "__main__":
    while True:
        mask_c = Mask_Client()
        cate = mask_c.mask_client()
        if cate is not None:
            cate = cate_map[cate]
            rospy.loginfo('Get Mask! The category is {}'.format(cate))
        else:
            rospy.loginfo('WARNING! No instance detected! Please check again!')
